'''
Descripttion: 阎老师自己做的数据集测试
Author: Haixu He
Date: 2021-12-26 14:55:42
'''
from matplotlib import markers
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.ticker import MaxNLocator
from models.bfast import BFAST
from datasets import *

plt.rc('font', family='Times New Roman')
plt.rcParams['font.size'] = 18
line_color = 'black'

def interp_nans(x):
    x_nn = np.array(pd.DataFrame(x).interpolate().values.ravel().tolist())
    nans = x_nn[np.isnan(x)]
    return x_nn, nans

def plot(name, y, x, f, season, level=0.05, h=0.15, max_iter=10, nan_clr="crimson"):
    def segmented_plot(arr, bp):
        prev = 0
        vals = np.concatenate((bp, [arr.shape[0] - 1]))
        for i, s in enumerate(vals + 1):
            ind = max(0, prev-1)
            ax.plot(x[ind:s], arr[ind:s], color=line_color)
            ax.axvline(x[ind], color=line_color, linestyle='--')
            prev = s

    def add_nans(y_n):
        if nans:
            ax.scatter(x_n, y_n, color=nan_clr, label="missing values", marker="x")
            ax.legend()

    def plot_data():
        # plot Y
        ax.set_ylabel('data')
        ax.set_xticks([])
        ax.plot(x, y, color=line_color)

        add_nans(y_n)
    def plot_sensonal():
        # plot sensonal
        ax.set_xticks([])
        ax.set_ylabel('seasonal')
        ax.yaxis.tick_right()
        if St_bp is not None:
            segmented_plot(St, St_bp)
        else:
            ax.plot(x, St, color=line_color)
        add_nans(St_n)
    def plot_trend():
        ax.set_xticks([])
        ax.set_ylabel('trend')
        if Tt_bp is not None:
            segmented_plot(Tt, Tt_bp)
        else:
            ax.plot(x, Tt, color=line_color)
        if season == "none":
            ax.yaxis.tick_right()
        add_nans(Tt_n)
    def plot_remainder():
        xx = x.flatten()
        ax.set_ylabel('remainder')
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        ax.axhline(0, color=line_color)
        if season != "none":
            ax.yaxis.tick_right()
        ax.bar(xx, Rt, color=line_color, width=(xx[1] - xx[0]) * 0.6)
        add_nans(Rt_n)

    bfast_output = BFAST(y, x, f, season=season, level=level, h=h, max_iter=max_iter)
    vo = bfast_output.output
    nans = np.isnan(y).any()

    x_n = x[np.isnan(y)]
    y, y_n = interp_nans(y)
    Tt, Tt_n = interp_nans(vo.trend)
    St, St_n = interp_nans(vo.season)
    Rt, Rt_n = interp_nans(vo.remainder)
    Tt_bp = vo.trend_breakpoints
    St_bp = vo.season_breakpoints

    figsz = (12, 8) if season != "none" else (12, 6)
    subfigs = 4 if season != "none" else 3

    fig = plt.figure(figsize=figsz)
    fig.suptitle(name, fontsize=24)


    ax = fig.add_subplot(subfigs, 1, 1)
    plot_data()   # 绘制data
    if season != 'none':
        ax = fig.add_subplot(subfigs, 1, 2)
        plot_sensonal()  # 绘制sensonal

    # plot trend
    ax = fig.add_subplot(subfigs, 1, subfigs-1)
    plot_trend()    # 绘制trend

    # plot remainder
    ax = fig.add_subplot(subfigs, 1, subfigs)
    plot_remainder()    # 绘制remainder

    fig.align_labels()   # 对齐标题
    plt.xlabel('Date')
    plt.subplots_adjust(hspace=0)
    plt.savefig('demo/{}.png'.format(name))
    # plt.show()

if __name__ == '__main__':
    plot("ndvi_simulat", ndvi_simulat, ndvi_simulat_dates, ndvi_simulat_freq, "harmonic")
    # plot("harvest", harvest, harvest_dates, harvest_freq, "harmonic")
    # plot("nile", nile, nile_dates, None, "none")
    # plot("SIMTS", simts_sum, simts_dates, simts_freq, "harmonic", level=0.35, h=0.3, max_iter=2)
    # plot("NDVI", ndvi, ndvi_dates, ndvi_freq, "dummy")
    # plot("nhtemp", nhtemp, nhtemp_dates, None, "none")
    # plot("uk_driver_deaths", uk_driver_deaths, uk_driver_deaths_dates, uk_driver_deaths_freq, "dummy")

